#include "Structures.fxh"
#include "Lighting.fxh"

cbuffer PerFrame {
	uniform float4x4 WVP; //WorldViewProjection
	uniform float4x4 World;
	uniform float4x4 WorldIT; //WorldInverseTranspose
	uniform float3 EyePos;
};

static const float SHADOW_EPSILON = 0.001f;
static const float SMAP_SIZE = 1024.0f;
static const float SMAP_DX = 1.0f / SMAP_SIZE;

//CBuffer for lighting values
cbuffer Lighting {
	uniform float3 GlobalAmbient : GLOBALAMBIENT = { 0.0f, 0.0f, 0.0f };
	uniform Light Lights[4] : LIGHTLIST;
	uniform int NumLights : LIGHTCOUNT = 0;
};

//CBuffer for material properties
cbuffer PerMaterial {
	uniform float3 MatDiffuse = { 1.0f, 1.0f, 1.0f};
	uniform float3 MatAmbient = { 0.2f, 0.2f, 0.2f};
	uniform float3 MatEmissive = { 0.0f, 0.0f, 0.0f};
	uniform float3 MatSpecular = { 0.0f, 0.0f, 0.0f};
	uniform float MatShininess = 16.0f;
	uniform float Alpha = 1.0f;
};

cbuffer Shadows {
  uniform float4x4 LightVP;
  uniform float DepthBias = 0.001f;
  uniform float3 LightDirection;
};

uniform Texture2D DiffuseMap;
uniform SamplerState DiffuseMapSampler {
	Filter = MIN_MAG_MIP_LINEAR;
	AddressU = Wrap;
	AddressV = Wrap;
};

uniform Texture2D ShadowMap;
uniform SamplerState ShadowMapSampler {
  Filter = MIN_MAG_MIP_POINT;
	AddressU = Wrap;
	AddressV = Wrap;
};

struct VSOutputNmTxShadow {
	float4 PositionPS : SV_POSITION;
	float3 PositionWS : TEXCOORD0;
	float3 Normal : TEXCOORD1;
	float2 TexCoord : TEXCOORD2;
	float4 TexCoordShadow : TEXCOORD3;
};

//
// **** LitBasicTexture ****
//

VSOutputNmTxShadow VS_LitBasicNmTx(VSInputNmTx input) {
	float4 pos = float4(input.Position, 1.0);

	VSOutputNmTxShadow output;
	output.PositionPS = mul(pos, WVP);
	output.PositionWS = mul(pos, World).xyz;
	output.Normal = normalize(mul(float4(input.Normal, 0.0), WorldIT).xyz);
	output.TexCoord = input.TexCoord;
	output.TexCoordShadow = mul(float4(output.PositionWS, 1.0f), LightVP);
	return output;
}

float CalcShadowFactor(float4 projTex) {
	projTex.xyz /= projTex.w;

	if(projTex.x < -1.0f || projTex.x > 1.0f || 
	   projTex.y < -1.0f || projTex.y > 1.0f || 
	   projTex.z < 0.0f) {
		return 0.0f;
	}

	//Transform from NDC space to texture space
	projTex.x = .5f * projTex.x + 0.5f;
	projTex.y = -.5f * projTex.y + 0.5f;

	float depth = projTex.z;

	float s0 = ShadowMap.Sample(ShadowMapSampler, projTex.xy).r;
	float s1 = ShadowMap.Sample(ShadowMapSampler, projTex.xy + float2(SMAP_DX, 0)).r;
	float s2 = ShadowMap.Sample(ShadowMapSampler, projTex.xy + float2(0, SMAP_DX)).r;
	float s3 = ShadowMap.Sample(ShadowMapSampler, projTex.xy + float2(SMAP_DX, SMAP_DX)).r;

	float result0 = depth <= s0 + SHADOW_EPSILON;
	float result1 = depth <= s1 + SHADOW_EPSILON;
	float result2 = depth <= s2 + SHADOW_EPSILON;
	float result3 = depth <= s3 + SHADOW_EPSILON;

	float2 texelPos = SMAP_SIZE * projTex.xy;

	float2 t = frac(texelPos);

	return lerp(lerp(result0, result1, t.x), lerp(result2, result3, t.x), t.y);
}

float4 PS_LitBasicNmTx(VSOutputNmTxShadow input) : SV_TARGET {
	//Individual light contribs from each light
	float3 lightContribs = float3(0,0,0);

	//Get the diffuse map color
	float4 albedo = DiffuseMap.Sample(DiffuseMapSampler, input.TexCoord);

	Material mat;
	mat.Ambient = MatAmbient;
	mat.Diffuse = MatDiffuse;
	mat.Specular = MatSpecular;
	mat.Shininess = MatShininess;

	//Calculate the V vector for specular highlights
	float3 V = normalize(EyePos - input.PositionWS);

	float3 normal = input.Normal;
	float3 pos = input.PositionWS;

	float shadowFactor = CalcShadowFactor(input.TexCoordShadow);

	for(int i = 0; i < NumLights; i++) {
		Light l = Lights[i];
		LightResult result = ComputeSingleLight(l, pos, normal, V, mat);
		float3 diffuse = (result.Diffuse + result.Ambient) * albedo.rgb;
		float3 specular = result.Specular;
		lightContribs += ComputeAttenuation(pos, l) * SpotEffect(pos, l) * (diffuse + specular) * shadowFactor;
	}

	float3 color = (MatAmbient * GlobalAmbient) + MatEmissive + lightContribs;

	return float4(color, Alpha * albedo.a);
}

technique10 LitBasicTexture {
	pass Pass0 {
		SetGeometryShader(0);
		SetVertexShader(CompileShader(vs_4_0, VS_LitBasicNmTx()));
		SetPixelShader(CompileShader(ps_4_0, PS_LitBasicNmTx()));
	}
}